import matplotlib.pyplot as plt
import numpy as np
from test_4_00 import create_polar_system

# 创建极坐标系
fig, ax = create_polar_system(r_range=(0, 5))

# 绘制双扭线 r^2 = a^2 * cos(2θ)
a = 2  # 双扭线参数

# 生成角度值
theta = np.linspace(0, 2*np.pi, 4000)

# 计算半径值
# r^2 = a^2 * cos(2θ)，所以 r = ±a * sqrt(cos(2θ))
# 只有当 cos(2θ) ≥ 0 时，r 才有实数解
r_squared = a**2 * np.cos(2*theta)
valid_indices = r_squared >= 0  # 找到有效的点（cos(2θ) ≥ 0）

# 计算正负两支
r_positive = np.sqrt(r_squared[valid_indices])
r_negative = -np.sqrt(r_squared[valid_indices])

# 绘制双扭线的两支
ax.plot(theta[valid_indices], r_positive, 'g-', linewidth=2, label='双扭线 r^2=(a^2)cos(2θ)')
ax.plot(theta[valid_indices], r_negative, 'g-', linewidth=2)

# 添加图例
ax.legend(loc='upper right', bbox_to_anchor=(1.15, 1.15))

# 标记极点
ax.plot(0, 0, 'ko', markersize=6, label='极点')
ax.text(0.1, 0.3, 'O', fontsize=10, color='black')

plt.tight_layout()
plt.show()
